from typing import TypedDict, List, Optional, Dict, Any
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage


class AgentState(TypedDict):
    """Agent状态定义"""
    # 用户输入
    user_input: str

    # 对话历史 (压缩形式)
    chat_history: List[BaseMessage]

    # 检索到的相关文档
    retrieved_docs: Optional[List[str]]

    # 工具调用结果
    tools_output: Optional[str]

    # 生成的响应
    response: Optional[str]

    # 当前处理步骤
    current_step: str

    # 会话ID
    session_id: str


def init_state(session_id: str, user_input: str) -> AgentState:
    """初始化对话状态"""
    return {
        "user_input": user_input,
        "chat_history": [],
        "retrieved_docs": None,
        "tools_output": None,
        "response": None,
        "current_step": "start",
        "session_id": session_id
    }


def update_history(state: AgentState, message: str, is_user=True):
    """更新对话历史"""
    if is_user:
        state["chat_history"].append(HumanMessage(content=message))
    else:
        state["chat_history"].append(AIMessage(content=message))
    return state